import torch
import sys

sys.path.append("/Users/zzq/Developer/python/SAM/reproduction/infini-transformer/")
from infini_transformer import InfiniTransformer

from infini_transformer import YaRNEmbeddings


# 实例化YaRNEmbeddings，设置其参数
embedder = YaRNEmbeddings(
    dim=64,  # 必须与InfiniTransformer中的dim_key参数匹配
    seq_len=2048,  # 必须与InfiniTransformer中的segment_len参数匹配
    context_len=32768,  # 上下文窗口长度
    context_len_ext=65536,  # 扩展的上下文窗口长度
    dim_embedding_pct=0.5,  # 使用的嵌入维度百分比
    base=10000,  # 基数
    alpha=1,  # 动态缩放的最小值
    beta=32,  # 动态缩放的最大值
    length_scale=None,  # 长度缩放因子
)

# 创建InfiniTransformer的实例，并设置其参数
tfm = InfiniTransformer(
    dim_input=768,  # 输入的维度
    dim_hidden=2048,  # 隐藏层的维度
    dim_key=64,  # 必须与YaRNEmbeddings中的dim参数匹配
    dim_value=64,  # 值的维度
    num_heads=8,  # 注意力头的数量
    activation="ffngeglu",  # 激活函数类型
    segment_len=2048,  # 必须与YaRNEmbeddings中的seq_len参数匹配
    update="delta",  # 更新策略
    causal=True,  # 是否为因果模型
    positional_embedder=embedder,  # 使用YaRNEmbeddings实例
    init_state_learnable=False,  # 初始状态是否可学习
    dropout=0.1,  # dropout率
)

# 创建一个随机张量，形状为(2, 65536, 768)
batch = torch.randn(2, 65536, 768)  # 批次大小  # 序列长度  # 输入维度

# 将批次数据传入InfiniTransformer实例，得到输出
output = tfm(batch)
print(output)
